#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/ktime.h>
#include <linux/limits.h>
#include <linux/sched.h>

#include "trap_io.h"

/**
 * extract32:
 * @value: the value to extract the bit field from
 * @start: the lowest bit in the bit field (numbered from 0)
 * @length: the length of the bit field
 *
 * Extract from the 32 bit input @value the bit field specified by the
 * @start and @length parameters, and return it. The bit field must
 * lie entirely within the 32 bit word. It is valid to request that
 * all 32 bits are returned (ie @length 32 and @start 0).
 *
 * Returns: the value of the bit field extracted from the input value.
 */
static inline uint32_t extract32(uint32_t value, int start, int length)
{
    return (value >> start) & (~0U >> (32 - length));
}

/**
 * sextract32:
 * @value: the value to extract the bit field from
 * @start: the lowest bit in the bit field (numbered from 0)
 * @length: the length of the bit field
 *
 * Extract from the 32 bit input @value the bit field specified by the
 * @start and @length parameters, and return it, sign extended to
 * an int32_t (ie with the most significant bit of the field propagated
 * to all the upper bits of the return value). The bit field must lie
 * entirely within the 32 bit word. It is valid to request that
 * all 32 bits are returned (ie @length 32 and @start 0).
 *
 * Returns: the sign extended value of the bit field extracted from the
 * input value.
 */
static inline int32_t sextract32(uint32_t value, int start, int length)
{
    /* Note that this implementation relies on right shift of signed
     * integers being an arithmetic shift.
     */
    return ((int32_t)(value << (32 - length - start))) >> (32 - length);
}
/*
 * Load/store (immediate post-indexed)
 * Load/store (immediate pre-indexed)
 * Load/store (unscaled immediate)
 *
 * 31 30 29   27  26 25 24 23 22 21  20    12 11 10 9    5 4    0
 * +----+-------+---+-----+-----+---+--------+-----+------+------+
 * |size| 1 1 1 | V | 0 0 | opc | 0 |  imm9  | idx |  Rn  |  Rt  |
 * +----+-------+---+-----+-----+---+--------+-----+------+------+
 *
 * idx = 01 -> post-indexed, 11 pre-indexed, 00 unscaled imm. (no writeback)
         10 -> unprivileged
 * V = 0 -> non-vector
 * size: 00 -> 8 bit, 01 -> 16 bit, 10 -> 32 bit, 11 -> 64bit
 * opc: 00 -> store, 01 -> loadu, 10 -> loads 64, 11 -> loads 32
 */
static int decode_ldst_reg_imm9(uint32_t insn,
                                int opc,
                                int size,
                                int rt,
                                bool is_vector,
                                mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;

    int rn = extract32(insn, 5, 5);
    int imm9 = sextract32(insn, 12, 9);
    int idx = extract32(insn, 10, 2);
    int opc0 = extract32(opc, 0, 1);
    bool is_signed = false;
    bool is_store = false;
    bool is_extended = false;
    bool is_unpriv = (idx == 2);
    bool post_index;
    bool writeback;
    int regsize;
    uint64_t address;

    if (is_vector || is_unpriv) {
        return ret;
    } else {
        if (size == 3 && opc == 2) {
            /* PRFM - prefetch */
            return ret;
        }
        if (opc == 3 && size > 1) {
            return ret;
        }
        is_store = (opc == 0);
        is_signed = extract32(opc, 1, 1);
        is_extended = (size < 3) && extract32(opc, 0, 1);
    }

    switch (idx) {
    case 0:
    case 2:
        post_index = false;
        writeback = false;
        break;
    case 1:
        post_index = true;
        writeback = true;
        break;
    case 3:
        post_index = false;
        writeback = true;
        break;
    default:
        return ret;
    }

    if (rn == 31) {
        /* Stack Point register - not support */
        return ret;
    }

    if (is_signed) {
        regsize = opc0 ? 32 : 64;
    } else {
        regsize = size == 3 ? 64 : 32;
    }

    address = get_x_reg(mmio->regs, rn);
    if (!post_index) {
        mmio->address = (void __iomem *)(address + (s64)imm9);
    } else {
        mmio->address = (void __iomem *)address;
    }


    mmio->regx = rt;
    mmio->is_store = is_store;
    mmio->access_size = regsize;
    mmio->sign_extend_read = is_signed &&  is_extended;

    if (is_store)
        mmio->value = get_x_reg(mmio->regs, rt);

    if (writeback) {
        set_x_reg(mmio->regs, rn, (unsigned long)(address + (s64)imm9));
    }

    ret = STATUS_OK;

    return ret;
}

/*
 * Load/store (unsigned immediate)
 *
 * 31 30 29   27  26 25 24 23 22 21        10 9     5
 * +----+-------+---+-----+-----+------------+-------+------+
 * |size| 1 1 1 | V | 0 1 | opc |   imm12    |  Rn   |  Rt  |
 * +----+-------+---+-----+-----+------------+-------+------+
 *
 * For non-vector:
 *   size: 00-> byte, 01 -> 16 bit, 10 -> 32bit, 11 -> 64bit
 *   opc: 00 -> store, 01 -> loadu, 10 -> loads 64, 11 -> loads 32
 * For vector:
 *   size is opc<1>:size<1:0> so 100 -> 128 bit; 110 and 111 unallocated
 *   opc<0>: 0 -> store, 1 -> load
 * Rn: base address register (inc SP)
 * Rt: target register
 */
static int decode_ldst_reg_unsigned_imm(uint32_t insn, int opc,
                            int size,
                            int rt,
                            bool is_vector,
                            mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;

    int rn = extract32(insn, 5, 5);
    unsigned int imm12 = extract32(insn, 10, 12);
    unsigned int offset;
    uint64_t address;
    int opc0 = extract32(opc, 0, 1);
    int regsize;

    bool is_store;
    bool is_signed = false;
    bool is_extended = false;

    if (is_vector) {
        return ret;
    } else {
        if (size == 3 && opc == 2) {
             /* PRFM - prefetch */
            return ret;
        }
        if (opc == 3 && size > 1) {
            /* unallocated */
            return ret;
        }
        is_store = (opc == 0);
        is_signed = extract32(opc, 1, 1);
        is_extended = (size < 3) && extract32(opc, 0, 1);
    }

    if (rn == 31) { 
        /* Stack Point register - not support */
        return ret;
    }

    address = get_x_reg(mmio->regs, rn);
    offset = imm12 << size;

    if (is_signed) {
        regsize = opc0 ? 32 : 64;
    } else {
        regsize = size == 3 ? 64 : 32;
    }

    mmio->regx = rt;
    mmio->address = (void __iomem *)(address + (uint64_t)offset);
    mmio->is_store = is_store;
    mmio->access_size = regsize;
    mmio->sign_extend_read = is_signed &&  is_extended;

    if (is_store)
        mmio->value = get_x_reg(mmio->regs, rt);

    ret = STATUS_OK;

    return ret;
}

#if 0
/*
 * Load/store (register offset)
 *
 * 31 30 29   27  26 25 24 23 22 21  20  16 15 13 12 11 10 9  5 4  0
 * +----+-------+---+-----+-----+---+------+-----+--+-----+----+----+
 * |size| 1 1 1 | V | 0 0 | opc | 1 |  Rm  | opt | S| 1 0 | Rn | Rt |
 * +----+-------+---+-----+-----+---+------+-----+--+-----+----+----+
 *
 * For non-vector:
 *   size: 00-> byte, 01 -> 16 bit, 10 -> 32bit, 11 -> 64bit
 *   opc: 00 -> store, 01 -> loadu, 10 -> loads 64, 11 -> loads 32
 * For vector:
 *   size is opc<1>:size<1:0> so 100 -> 128 bit; 110 and 111 unallocated
 *   opc<0>: 0 -> store, 1 -> load
 * V: 1 -> vector/simd
 * opt: extend encoding (see DecodeRegExtend)
 * S: if S=1 then scale (essentially index by sizeof(size))
 * Rt: register to transfer into/out of
 * Rn: address register or SP for base
 * Rm: offset register or ZR for offset
 */
static int decode_ldst_reg_roffset(uint32_t insn,
                                   int opc,
                                   int size,
                                   int rt,
                                   bool is_vector,
                                   mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;
    int rn = extract32(insn, 5, 5);
    int shift = extract32(insn, 12, 1);
    int rm = extract32(insn, 16, 5);
    int opt = extract32(insn, 13, 3);
    bool is_signed = false;
    bool is_store = false;
    bool is_extended = false;
    int opc0 = extract32(opc, 0, 1);
    int regsize;
    uint64_t address;
    unsigned int offset;

    if (extract32(opt, 1, 1) == 0) {
        return ret;
    }

    if (is_vector) {
        return ret;
    } else {
        if (size == 3 && opc == 2) {
            /* PRFM - prefetch */
            return ret;
        }
        if (opc == 3 && size > 1) {
            return ret;
        }
        is_store = (opc == 0);
        is_signed = extract32(opc, 1, 1);
        is_extended = (size < 3) && extract32(opc, 0, 1);
    }

    if (rn == 31) {
        /* Stack Point register - not support */
        return ret;
    }

    if (is_signed) {
        regsize = opc0 ? 32 : 64;
    } else {
        regsize = size == 3 ? 64 : 32;
    }

    address = get_x_reg(mmio->regs, rn);
    offset = get_x_reg(mmio->regs, rm);

    return ret;
}
#endif

static int decode_ldst_reg(uint32_t insn, mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;

    int rt = extract32(insn, 0, 5);
    int opc = extract32(insn, 22, 2);
    bool is_vector = extract32(insn, 26, 1);
    int size = extract32(insn, 30, 2);

    if (is_vector)
        return ret;

    switch (extract32(insn, 24, 2)) {
    case 0:
        if (extract32(insn, 21, 1) == 0) {
            /* Load/store register (unscaled immediate)
             * Load/store immediate pre/post-indexed
             * Load/store register unprivileged
             */
            ret = decode_ldst_reg_imm9(insn, opc, size, rt, is_vector, mmio);
            return ret;
        }
        switch (extract32(insn, 10, 2)) {
        case 2:
            //ret = decode_ldst_reg_roffset(insn, opc, size, rt, is_vector, mmio);
            return ret;
        case 0: /* load and store atomic*/
        default:
            return ret;
        }
    case 1:
        ret = decode_ldst_reg_unsigned_imm(insn, opc, size, rt, is_vector, mmio);
        return ret;
    }

    return ret;
}

static int decode_ldst(uint32_t insn, mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;

    switch (extract32(insn, 24, 6)) {
    case 0x38: case 0x39:
    case 0x3c: case 0x3d: /* Load/store register (all forms) */
        ret = decode_ldst_reg(insn, mmio);
        break;
    case 0x08: /* Load/store exclusive */
    case 0x18: case 0x1c: /* Load register (literal) */
    case 0x28: case 0x29: case 0x2c: case 0x2d: /* Load/store pair (all forms) */
    case 0x0c: /* AdvSIMD load/store multiple structures */
    case 0x0d: /* AdvSIMD load/store single structure */
    default:
        break;
    }

    return ret;
}

/* C3.1 A64 instruction index by encoding */
int decode_aarch64_insn(uint32_t insn, mmio_op_t *mmio)
{
    int ret = STATUS_FAIL;

    switch (extract32(insn, 25, 4)) {
    case 0x4:
    case 0x6:
    case 0xc:
    case 0xe:      /* Loads and stores */
        ret = decode_ldst(insn, mmio);
        break;

    case 0x0: case 0x1: case 0x2: case 0x3: /* UNALLOCATED */
    case 0x8: case 0x9: /* Data processing - immediate */
    case 0xa: case 0xb: /* Branch, exception generation and system insns */
    case 0x5: case 0xd:      /* Data processing - register */
    case 0x7: case 0xf:      /* Data processing - SIMD and floating point */
    default:
        break;
    }

    return ret;
}
